package edu.washington.ling.syntric.pipeline;

import java.io.IOException;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;

import org.apache.commons.configuration.ConfigurationException;
import org.apache.commons.configuration.XMLConfiguration;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.config.Configuration;
import edu.washington.ling.syntric.fold.Fold;
import edu.washington.ling.syntric.fold.FoldConfiguration;
import edu.washington.ling.syntric.io.ExampleDoc;
import edu.washington.ling.syntric.io.ExampleDocReader;
import edu.washington.ling.syntric.io.Example;
import edu.washington.ling.syntric.io.TreeReader;
import edu.washington.ling.syntric.io.TreeRepresentation;
import edu.washington.ling.syntric.proc.ExtractBinaryRules;
import edu.washington.ling.syntric.proc.ExtractUnaryRules;
import edu.washington.ling.syntric.struct.error.PrintingExceptionHandler;
import edu.washington.ling.syntric.struct.error.TreeMergeException;
import edu.washington.ling.syntric.struct.error.UnreadableRepresentationException;
import edu.washington.ling.syntric.struct.rule.BiDSPattern;
import edu.washington.ling.syntric.struct.rule.BiPSPattern;
import edu.washington.ling.syntric.struct.rule.BiRule;
import edu.washington.ling.syntric.struct.rule.UniDSPattern;
import edu.washington.ling.syntric.struct.rule.UniPSPattern;
import edu.washington.ling.syntric.struct.rule.UniRule;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.ConvertUtils;
import edu.washington.ling.syntric.util.MergeUtils;

class TrainThread extends Thread {
	private PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	private List<String> ids;
	private ExampleDoc exf;
	private String outputPrefix;
	private Configuration config;
	
	public TrainThread(List<String> ids, ExampleDoc exf, String outputPrefix, Configuration config) {
		this.ids = ids;
		this.exf = exf;
		this.outputPrefix = outputPrefix;
		this.config = config;
	}
	private Map<BiDSPattern, List<BiPSPattern>> getBiRuleMapping(Set<BiRule> rules) {
		Map<BiDSPattern, List<BiPSPattern>> ruleMap = new HashMap<BiDSPattern, List<BiPSPattern>>();
		for (BiRule rule: rules) {
			if (!ruleMap.containsKey(rule.getDSPattern())) {
				ruleMap.put(rule.getDSPattern(), new LinkedList<BiPSPattern>());
			}
			ruleMap.get(rule.getDSPattern()).add(rule.getPSPattern());
		}
		return ruleMap;
	}
	private Map<UniDSPattern, List<UniPSPattern>> getUniRuleMapping(Set<UniRule> rules) {
		Map<UniDSPattern, List<UniPSPattern>> ruleMap = new HashMap<UniDSPattern, List<UniPSPattern>>();
		for (UniRule rule: rules) {
			if (!ruleMap.containsKey(rule.getDSPattern())) {
				ruleMap.put(rule.getDSPattern(), new LinkedList<UniPSPattern>());
			}
			ruleMap.get(rule.getDSPattern()).add(rule.getPSPattern());
		}
		return ruleMap;
	}
	private void train() {
		Map<BiRule, List<Example>> biRuleExMap = new HashMap<BiRule, List<Example>>();
		Map<UniRule, List<Example>> uniRuleExMap = new HashMap<UniRule, List<Example>>();
		for (String id: ids) {
			try {
				Example exr = exf.getExampleWithID(id);
				TreeRepresentation dsRep = exr.getTreeRep("DS-Const");
				TreeRepresentation psRep = exr.getTreeRep("PS");
				TreeReader dsr = TreeReader.newReader(dsRep.getFormat());
				TreeReader psr = TreeReader.newReader(psRep.getFormat());
				try {
					Tree ds = dsr.readTree(dsRep);
					Tree ps = psr.readTree(psRep);
					Tree tree = MergeUtils.mergeDSandPSTrees(ds, ps);
					Set<UniRule> uniRules = ExtractUnaryRules.extractUnaryRules(tree, config);
					for(UniRule uniRule: uniRules) {
						if (!uniRuleExMap.containsKey(uniRule)) {
							uniRuleExMap.put(uniRule, new LinkedList<Example>());
						}	
						uniRuleExMap.get(uniRule).add(exr);
					}
					Set<BiRule> biRules = ExtractBinaryRules.extractBinaryRules(tree, config);
					for(BiRule biRule: biRules) {
						if (!biRuleExMap.containsKey(biRule)) {
							biRuleExMap.put(biRule, new LinkedList<Example>());
						}	
						biRuleExMap.get(biRule).add(exr);
					}
				}
				catch (TreeMergeException e) {
					String errStr = String.format("Could not merge representations for example '%s'", id);
					pHandler.handle(e, errStr);
				}
				catch (UnreadableRepresentationException e) {
					String errStr = String.format("Could not process representation for example '%s'", id);
					pHandler.handle(e, errStr);
				}
			}
			catch(Exception e) {
				pHandler.handle(e, "Training thread failed.", System.out);
				e.printStackTrace(System.out);
				System.exit(-1);
			}
		}
		Map<UniDSPattern, List<UniPSPattern>> uniRuleMap = getUniRuleMapping(uniRuleExMap.keySet());
		Map<BiDSPattern, List<BiPSPattern>> biRuleMap = getBiRuleMapping(biRuleExMap.keySet());
		ExtractUnaryRules.printToFile(uniRuleMap, uniRuleExMap, outputPrefix + "-uni-rules.txt", "text");
		ExtractBinaryRules.printToFile(biRuleMap, biRuleExMap, outputPrefix + "-bi-rules.txt", "text");
	}
	
	@Override
	public void run() {
		train();
	}
}

public class Train {
	private static TrainOptions opts = new TrainOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	public static List<String> getCompatibleExampleIDs(ExampleDoc exf) {
		List<String> compatExampleIDs = new LinkedList<String>();
		for (String exampleID: exf.getIDs()) {
			if (ConvertUtils.isCompatible(exf.getExampleWithID(exampleID))) {
				compatExampleIDs.add(exampleID);
			}
		}
		return compatExampleIDs;
	}
	
	public static void train(ExampleDoc exf, String outputPrefix, 
			Configuration config, FoldConfiguration foldConfig, int maxThreads) {
		
		MultithreadedProcedure mtProc = new MultithreadedProcedure(maxThreads);
		
		if (foldConfig == null) {
			List<String> compatExampleIDs = getCompatibleExampleIDs(exf);
			mtProc.addThread(new TrainThread(compatExampleIDs, exf, outputPrefix, config));
		} 
		else {
			for (int i = 0; i < foldConfig.numFolds(); i++) {
				Fold fold = foldConfig.getFold(i);
				String foldPrefix = String.format("%s-fold%d", outputPrefix, i);
				mtProc.addThread(new TrainThread(fold.getTrainExampleIDs(), exf, foldPrefix, config));
			}
		}
		mtProc.run();
		System.out.println("Training was successful.");
	}
	
	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			Configuration config;
			if (opts.configFilePath != null) {
				XMLConfiguration xmlConfig = new XMLConfiguration(opts.configFilePath);
				config = Configuration.init(xmlConfig);
			} else {
				config = Configuration.initDefault();
			}
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			FoldConfiguration foldConfig = null;
			if (opts.existingFoldConfigFile != null) {
				List<String> compatExampleIDs = getCompatibleExampleIDs(exf);
				foldConfig = FoldConfiguration.init(opts.existingFoldConfigFile);
				if (!foldConfig.checkExampleIDs(compatExampleIDs)) {
					System.out.println("Error: fold configuration example IDs " +
							"do not match the input file");
					System.exit(-1);
				}
			}
			if (opts.newFoldConfigFile != null) {
				List<String> compatExampleIDs = getCompatibleExampleIDs(exf);
				foldConfig = FoldConfiguration.init(compatExampleIDs, opts.numFolds);
				foldConfig.writeXMLFile(opts.newFoldConfigFile);
			}
			train(exf, opts.outputPrefix, config, foldConfig, opts.maxThreads);
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("Usage: Train");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(-1);
		}
		catch(IOException e) {
			String errStr = String.format("Example file '%s' could not be read", opts.inputFilePath);
			pHandler.handle(e, errStr, System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
		catch(ConfigurationException e) {
			String errStr = String.format("Configuration file '%s' could not be read", opts.configFilePath);
			pHandler.handle(e, errStr, System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
		catch(NoSuchElementException e) {
			String errStr = String.format("Configuration file '%s' is incomplete", opts.configFilePath);
			pHandler.handle(e, errStr, System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		} 
		catch (IllegalArgumentException e) {
			String errStr = String.format("Configuration file '%s' has illegal values", opts.configFilePath);
			pHandler.handle(e, errStr, System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}
}
